#include <set>

#include "Rotor.hpp"
#include "Reflector.hpp"

/* advances Rotor 1, maybe advance Rotor 2 and 3 too */
void advance(Rotor &r1, Rotor &r2, Rotor &r3) {
	bool result = r1.advance();
	if (result) {
		result = r2.advance();
		if (result) {
			r3.advance();
		}
	}
}

char encodeDecode(char c, Rotor &r1, Rotor &r2, Rotor &r3, const Reflector &rr) {
	return r1.decrypt(r2.decrypt(r3.decrypt(rr.reflect(r3.encrypt(r2.encrypt(r1.encrypt(c)))))));
}

std::string encodeDecode(const std::string &plain, Rotor &r1, Rotor &r2, Rotor &r3, const Reflector &rr) {
	//copy the string
	std::string result = std::string(plain);
	//for each char
	for (unsigned int i = 0; i < plain.length(); i++) {
		//encode it
		result[i] = encodeDecode(plain[i], r1, r2, r3, rr);
		//and advance the rotors
		advance(r1, r2, r3);
	}
	return result;
}

std::set<std::string> bruteForce(const std::string &ciphertext, const std::string &knownPlaintext, Rotor &r1, Rotor &r2, Rotor &r3, const Reflector &rr) {
	std::set<std::string> result;

	//all possible rotor permitations
	int index = 0;
	Rotor* rotors[] = { &r1, &r2, &r3,
	                    &r1, &r3, &r2,
	                    &r2, &r1, &r3,
	                    &r2, &r3, &r1,
	                    &r3, &r1, &r2,
	                    &r3, &r2, &r1, };
	std::string configs[] = { "123",
	                          "132",
	                          "213",
	                          "231",
	                          "312",
	                          "321", };

	//iterate over all the permutations
	while (index != 18) {
		Rotor rotor1 = *rotors[index++];
		Rotor rotor2 = *rotors[index++];
		Rotor rotor3 = *rotors[index++];

		//the initial configurations
		char ch1 = 'A';
		char ch2 = 'A';
		char ch3 = 'A';
		
		//try all possible combinations
		while (true) {
			//initialize the rotors
			rotor1.reset(ch1);
			rotor2.reset(ch2);
			rotor3.reset(ch3);
			
			//decrypt and look for the plaintext
			std::string plaintext = encodeDecode(ciphertext, rotor1, rotor2, rotor3, rr);
			if (plaintext.find(knownPlaintext) != std::string::npos) {
				//and add the used config to the set of results
				std::string temp = std::string(configs[(index / 3) - 1]);
				temp += ch1;
				temp += ch2;
				temp += ch3;
				result.insert(temp);
			}
			
			//continue with the next combination
			ch1++;
			if (ch1 > 'Z') {
				ch1 = 'A';
				ch2++;
				if (ch2 > 'Z') {
					ch2 = 'A';
					ch3++;
					if (ch3 > 'Z') {
						break;
					}
				}
			}
		}
	}

	return result;
}
